

#include "glodef.h"
#include "errlog.h"
#include "grid3d.h"
#include "fe.h"
#include "fe_matr.h"
#include "vecs.h"
#include "ns_csplit.h"

#include "f3io.h"
#include "soah_common.h"

#include "glovar.h"

#include "ns_csplit.h"
/*
 #include "ns_pcorr.h"
 */

#include <stdlib.h>
#include <string.h>
#include <stdio.h>
// #include <time.h>

double current_solution_norm[];
double time_inf_solution_norm[];


char nrms_fname[FILENAME_MAX]="norms.dat";
vector eu, ev;
scalar ep, eq;

static void state_advection(void);
static void solution_norms(void);
static void print_inf_norms(void);

/****************************************************************************/

void init_state_solver(void)
{
    solver_data=InitCSplit(elemU, elemP);
    if(elemU==fe_bub_vec)
    {
        eu=alloc_vec(fe_ndof(grid, fe_bub));
        ev=alloc_vec(fe_ndof(grid, fe_bub));
    }
    else
    {
        eu=alloc_vec(fe_ndof(grid, elemU));
        ev=alloc_vec(fe_ndof(grid, elemU));
    }
    ep=alloc_scal(fe_ndof(grid, elemP));
    eq=alloc_scal(fe_ndof(grid, elemP));
    /*
     if(dump_bin > 0)
        sprintf(nrms_fname, "%s/norms.dat", state_dir());
    else
     */

    if(flags.shape_opt)
    {
        reference_nodes=calloc(grid->NN, sizeof(node_t));
        if(reference_nodes==NULL)
            error_msg("%s: Memory allocation error.", __func__);
        set_reference_nodes(grid->nodes);
        can_move_nodes=malloc(sizeof(int)*grid->NN);
        if(can_move_nodes==NULL)
            error_msg("%s: Memory allocation error.", __func__);
        for(int i=0; i<grid->NN; i++)
            if(reference_nodes[i].on_bdry)
            {
                // reference_nodes[i].bcond=BC_DIR;
                can_move_nodes[i]= -1;
            }
            else
                can_move_nodes[i]=1;

        /* local_move_grid_matrix is in soah_mesh.c */
        extern void local_move_grid_matrix(const grid3d_t *grid, int el,
                                           unsigned type, double ***S);
        move_grid_matr=fe_matr_alloc(fe_p1_vec, local_move_grid_matrix,
                                     eo_bc_dir, grid, TRUE, NULL);
        fe_matr_assemble(grid, 1, &move_grid_matr);
    }

    sprintf(nrms_fname, "norms.dat");
    atexit(done_state_solver);
}

/****************************************************************************/

void done_state_solver(void)
{
    DoneCSplit(solver_data);
    free_vec(eu);
    free_vec(ev);
    free_scal(ep);
    free_scal(eq);
    if(reference_nodes)
        free(reference_nodes);
    if(move_grid_matr)
        fe_matr_free(move_grid_matr);
}

/****************************************************************************/

int solve_state_ivp(vector u_init, scalar p_init)
{
    int ret=TRUE;

    double cost, old_cost;

    extern void test_shape_gradient(void);

    curr_ti=0;
    move_grid();

    PrintLogLine(78, '=');
    PrintLog("Time %-10f\n", (double) TIME);

    /* make initial pressure compatible with u_init */
#if defined(CS_INITIAL_PRESSURE)
    {
        int i, N=get_N_vec(u_init);
        vector ut_bc=alloc_vec(N);
        double *ut_bcX=get_X_vec(ut_bc),
                *ut_bcY=get_Y_vec(ut_bc),
                *ut_bcZ=get_Z_vec(ut_bc);
        double *u_initX=get_X_vec(u_init),
                *u_initY=get_Y_vec(u_init),
                *u_initZ=get_Z_vec(u_init);

        curr_ti=1;
        if(flags.moving_grid)
            move_grid();

        apply_bc_dir_vector(grid, elemU, TIME, bdry_vel, ut_bc);

        if(on_ival|on_oval)
            div_free_bc(grid, elemU, on_ival|on_oval, ut_bc);

        for(i=0; i<N; i++)
            if(fe_dof_bcond(grid, elemU, i)==BC_DIR)
            {
                ut_bcX[i]=(ut_bcX[i]-u_initX[i])/TStep;
                ut_bcY[i]=(ut_bcY[i]-u_initY[i])/TStep;
                ut_bcZ[i]=(ut_bcZ[i]-u_initZ[i])/TStep;
            }

        curr_ti=0;
        if(flags.moving_grid)
            move_grid();
        InitialPressure_cs(solver_data, u_init, ut_bc, p_init);
        make_scal_mean_0(grid, elemP, p_init);
        free_vec(ut_bc);
    }
#endif

    clean_glovecs();

    PRINTF("\rT = %6.3f; ", (double) (TIME));

    copy_vec(u_init, u);
    copy_scal(p_init, p);
    memset(time_inf_solution_norm, 0, sizeof(time_inf_solution_norm));
    solution_norms();

    total_cost=0.0;
    cost=CostFunctional();

    if(dump_bin>0)
    {
        create_dir(state_dir());
        save_state_time(state_dir());
    }
    for(curr_ti++; curr_ti<=total_ti; curr_ti++)
    {
        double tm=MPI_Wtime();

        PrintLogLine(78, '=');
        PrintLog("Time %-10f\n", (double) TIME);

        if(!flags.moving_grid)
        {
            copy_vec(u, adv_u);
        }
        else
        {
            move_grid();
            state_advection(); /* fills adv_u appropriately
             + transfers pressure */
        }
        copy_vec(u0, u1);
        copy_vec(u, u0);
        copy_scal(p0, p1);
        copy_scal(p, p0);

        /* do a timestep */
        ret = TimeStep_cs(solver_data, u, p,
                    u0, p0, u1, p1, flags.moving_grid||curr_ti==1);
        if(!ret)
            break;
        /* the solver may set p to 0 in one point.  To compute l2 norms we
         * need p to have mean 0, so
         */
        make_scal_mean_0(grid, elemP, p);

        if((dump_bin>0)&&(curr_ti%dump_bin)==0)
            save_state_time(state_dir());

        PRINTF("\rT = %6.3f; ", (double) (TIME));

        solution_norms();

        old_cost=cost;
        cost=CostFunctional();
#if (COST_FUNC_SQUARE == 0)
        total_cost+=TStep*0.5*(cost*cost+old_cost*old_cost);
#else
        total_cost+=TStep*0.5*(cost+old_cost);
#endif

        /*
         test_shape_gradient();
         */

        FFLUSH(stdout);

        PrintLog("CPU time for this timestep: %-6.3fsec.\n", (MPI_Wtime()-tm));
    }

#if (COST_FUNC_SQUARE == 0)
    total_cost=sqrt(total_cost);
#endif
    print_inf_norms();

    return ret;
}

/****************************************************************************/
/****************************************************************************/
/****************************************************************************/

#include "transfer.h"

static void state_advection(void)
{
    switch(flags.advection)
    {
        case adv_none:
            return;
        case adv_explicit:
        case adv_rk45:
        case adv_semi:
        {
            // vector ndisp=get_ndisp(alegrid); /* call after `move_grid()' */
            /* ATTENTION:  This works only in the particular case where
             * the advecting velocity `adv_u' is in fe_bub and the grid
             * displacement `ndisp' is in fe_p1.  If this is not the case,
             * then the injection/projection of `ndisp' to the space of
             * velocities must be done.
             */
            // if(elemU!=fe_bub||get_N_vec(ndisp)!=grid->NN)
            //    error_msg("ALE is only available for velocity in P1-bubble.\n");

            if(flags.use_ale)
            {
                error_msg("ALE is not available at this time\n");
                copy_vec(u, adv_u);
                /* for(int i=0; i<grid->NN; i++)
                {
                    adv_u->X[i]-=ndisp->X[i]/TStep;
                    adv_u->Y[i]-=ndisp->Y[i]/TStep;
                    adv_u->Z[i]-=ndisp->Z[i]/TStep;
                } */
                //axpy_vec(adv_u, -1.0/TStep, ndisp);
            }

            if(flags.transfer==transfer_taylor)
            {
                // transfer_taylor_scalar(ndisp, elemP, elemU);
            }
            else if(flags.transfer==transfer_interpolation)
            {
                error_msg("Transfer by interpolation is not implemented.\n");
            }

            if(!flags.use_ale)
                copy_vec(u, adv_u);
        }
            break;
        case adv_char:
            error_msg("Characteristics is not available at this time.\n");
            /* {
                double AGQ;

                use_vector(u)
                PrintLog("Method of characteristics.\n");
                for(int i=0; i<grid->NN; i++)
                {
                    adv_nodes[i].coords[0]=
                            alegrid->old_nodes[i].coords[0]+TStep*uX[i];
                    adv_nodes[i].coords[1]=
                            alegrid->old_nodes[i].coords[1]+TStep*uY[i];
                    adv_nodes[i].coords[2]=
                            alegrid->old_nodes[i].coords[2]+TStep*uZ[i];
                }
                done_vector(u)

                AGQ=GridQuality(grid, adv_nodes);
                if(AGQ<SMALLF)
                    error_msg("AGO < SMALLF");
                transfer_interpolation_vector(grid, elemU,
                                              adv_nodes, u, grid->nodes, u0);
                copy_vec(u0, u);
            }*/
            break;
    }
}

/****************************************************************************/

static void solution_norms(void)
{
    /* if exct_vel is NULL then eu is set to 0, same for exct_p and ep */
    interpolate_vector_lagr(grid, elemU, TIME, exct_vel, eu);
    interpolate_scalar_lagr(grid, elemP, TIME, exct_prs, ep);
    make_scal_mean_0(grid, elemP, ep);

    current_solution_norm[NORM_DUDT]=
            norm_vec_diff_l2(grid, elemU, u, u0)/TStep;

    current_solution_norm[NORM_ERR_U_L2]=norm_vec_diff_l2(grid, elemU, u, eu);
    current_solution_norm[NORM_ERR_U_H1]=norm_vec_diff_h1(grid, elemU, u, eu);
    current_solution_norm[NORM_ERR_P_L2]=norm_scal_diff_l2(grid, elemP, p, ep);
    current_solution_norm[NORM_ERR_P_H1]=norm_scal_diff_h1(grid, elemP, p, ep);

    current_solution_norm[NORM_U_L2]=norm_vec_diff_l2(grid, elemU, u, NULL);
    current_solution_norm[NORM_P_L2]=norm_scal_diff_l2(grid, elemP, p, NULL);
    current_solution_norm[NORM_GRAD_U_L2]=norm_grad_vec(grid, elemU, u);

    current_solution_norm[NORM_EXCT_U_L2]
            =norm_vec_diff_l2(grid, elemU, eu, NULL);
    current_solution_norm[NORM_EXCT_P_L2]
            =norm_scal_diff_l2(grid, elemP, ep, NULL);
    current_solution_norm[NORM_EXCT_GRAD_U_L2]=norm_grad_vec(grid, elemU, eu);

    PrintLogNorm("u-eu", "L2", current_solution_norm[NORM_ERR_U_L2]);
    PrintLogNorm("p-ep", "L2", current_solution_norm[NORM_ERR_P_L2]);
    PrintLogNorm("u-eu", "H1", current_solution_norm[NORM_ERR_U_H1]);
    PrintLogNorm("p-ep", "H1", current_solution_norm[NORM_ERR_P_H1]);
    PrintLogNorm("eu", "L2", current_solution_norm[NORM_EXCT_U_L2]);
    PrintLogNorm("ep", "L2", current_solution_norm[NORM_EXCT_P_L2]);
    PrintLogNorm("grad eu", "L2", current_solution_norm[NORM_EXCT_GRAD_U_L2]);
    PrintLogNorm("u", "L2", current_solution_norm[NORM_U_L2]);
    PrintLogNorm("p", "L2", current_solution_norm[NORM_P_L2]);
    PrintLogNorm("grad u", "L2", current_solution_norm[NORM_GRAD_U_L2]);
    PrintLogNorm("du/dt", "L2", current_solution_norm[NORM_DUDT]);

    if(curr_ti>=flags.time_order)
        for(int i=0; i<NORM_NUMBER; i++)
            time_inf_solution_norm[i]=fmax(time_inf_solution_norm[i],
                                           current_solution_norm[i]);

    if(exct_vel!=NULL)
    {
        PRINTF("||u-eu||_L2 = %10.4e", current_solution_norm[NORM_ERR_U_L2]);
    }
    else
    {
        if(flags.stationary)
        {
            PRINTF("||du/dt||_L2 = %10.4e", current_solution_norm[NORM_DUDT]);
        }
        else
        {
            PRINTF("||  u ||_L2 = %10.4e", current_solution_norm[NORM_U_L2]);
        }
    }
}

/****************************************************************************/

static void print_inf_norms(void)
{
    PrintLogLine(1, ' ');
    PrintLogLine(78, '=');
    PrintLogNorm("u-eu", "L2Linf", time_inf_solution_norm[NORM_ERR_U_L2]);
    PrintLogNorm("p-ep", "L2Linf", time_inf_solution_norm[NORM_ERR_P_L2]);
    PrintLogNorm("u-eu", "H1Linf", time_inf_solution_norm[NORM_ERR_U_H1]);
    PrintLogNorm("p-ep", "H1Linf", time_inf_solution_norm[NORM_ERR_P_H1]);
    PrintLogNorm("eu", "L2Linf", time_inf_solution_norm[NORM_EXCT_U_L2]);
    PrintLogNorm("ep", "L2Linf", time_inf_solution_norm[NORM_EXCT_P_L2]);
    PrintLogNorm("grad eu", "L2Linf",
                 time_inf_solution_norm[NORM_EXCT_GRAD_U_L2]);
    PrintLogNorm("u", "L2Linf", time_inf_solution_norm[NORM_U_L2]);
    PrintLogNorm("p", "L2Linf", time_inf_solution_norm[NORM_P_L2]);
    PrintLogNorm("grad u", "L2Linf", time_inf_solution_norm[NORM_GRAD_U_L2]);
    PrintLogNorm("du/dt", "L2Linf", time_inf_solution_norm[NORM_DUDT]);
    PrintLog("       %10s   %-20s = %.8g\n", "total_cost", "", total_cost);
    PRINTF("\r       %10s   %-20s = %.8g\n", "total_cost", "", total_cost);
}

void print_conv_norms(void)
{
    if(mpi_my_id!=0)
        return;
    FILE *F=fopen(nrms_fname, "a+t");
    if(F==NULL)
        return;
    fseek(F, 0, SEEK_SET);
    if(fgetc(F)==EOF)
    {
        fprintf(F, "%-12s    %-20s    %-14s  %-14s  %-14s  %-14s\n",
                "dT", "grid", "||u-eu||L2linf", "||p-ep||L2linf",
                "||u-eu||H1linf", "||p-ep||H1linf");
    }
    else
        fseek(F, 0, SEEK_END);
    /* fprintf(F, "%-12g    %-12g", (double) TStep, (double) GStep); */
    fprintf(F, "%-12g    %-20.20s", (double) TStep, grid->gridname);
    fprintf(F, "    %12e", time_inf_solution_norm[NORM_ERR_U_L2]);
    fprintf(F, "    %12e", time_inf_solution_norm[NORM_ERR_P_L2]);
    fprintf(F, "    %12e", time_inf_solution_norm[NORM_ERR_U_H1]);
    fprintf(F, "    %12e", time_inf_solution_norm[NORM_ERR_P_H1]);
    fprintf(F, "\n");
    fclose(F);
}

